#include "siamRpnTracker.h"
#include "elog.h"
#include "elog_file.h"
#include "dwarf_utils.h"
#include "timer_utils.h"

#include <sys/ioctl.h>

#include <opencv2/opencv.hpp>
#include <string>

#define DMA_BUF_SYNC_READ      (1 << 0)
#define DMA_BUF_SYNC_WRITE     (2 << 0)
#define DMA_BUF_SYNC_RW        (DMA_BUF_SYNC_READ | DMA_BUF_SYNC_WRITE)
#define DMA_BUF_SYNC_START     (0 << 2)
#define DMA_BUF_SYNC_END       (1 << 2)

typedef unsigned long long __u64;

struct dma_buf_sync {
	__u64 flags;
};

#define DMA_BUF_BASE		'b'
#define DMA_BUF_IOCTL_SYNC	_IOW(DMA_BUF_BASE, 0, struct dma_buf_sync)

int dma_sync_device_to_cpu(int fd) {
    struct dma_buf_sync sync = {0};

    sync.flags = DMA_BUF_SYNC_START | DMA_BUF_SYNC_RW;
    return ioctl(fd, DMA_BUF_IOCTL_SYNC, &sync);
}

int dma_sync_cpu_to_device(int fd) {
    struct dma_buf_sync sync = {0};

    sync.flags = DMA_BUF_SYNC_END | DMA_BUF_SYNC_RW;
    return ioctl(fd, DMA_BUF_IOCTL_SYNC, &sync);
}

/**
 * @brief 
 *      <static function> select and copy ROI of the original image to new buffer
 *  
 * @param dst_data
 *      address of new buffer
 * @param src_data 
 *      address of original image buffer
 * @param h :
 *      height of original image
 * @param w :
 *      width of original image
 * @param c :
 *      channel number of original image
 * @param x :
 *      ROI start point x
 * @param y :
 *      ROI start point y
 * @param size_x : 
 *      ROI width
 * @param size_y :
 *      ROI height
 * @param up_pad :
 *      up side padding size
 * @param buttom_pad : 
 *      buttom side padding size
 * @param left_pad :
 *      left side padding size
 * @param right_pad :
 *      right side padding size
 * 
 */
static void fast_copy_with_pad(unsigned char* dst_data, const unsigned char* src_data,  int h, int w, int c, int x, int y, int size_x, int size_y, int up_pad, int buttom_pad, int left_pad, int right_pad)
{
	int i, j, k;
	int dst_w = size_x + left_pad + right_pad;
	int dst_start_h = up_pad, dst_start_w = left_pad;
	for (i = 0; i < size_y; i++) {
		for (j = 0; j < size_x; j++) {
			for (k = 0; k < c; k++) {
				*(dst_data + (i + dst_start_h) * dst_w * c + (j + dst_start_w) * c + k) = *(src_data + ((i + y) * w) * c + (j + x) * c + k);
			}
		}
	}
}

/**
 * @brief 
 *      <static function> transpose a 3D image array to CHW format
 * @param dst_data 
 *      address of new buffer
 * @param src_data 
 *      address of original buffer
 * @param h 
 *      height of original image
 * @param w 
 *      width of original image
 * @param c 
 *      channel number of original image
 */
static void fast_transpose_201(unsigned char* dst_data, const unsigned char* src_data, int h, int w, int c) {
	int i, j, k;
	for (i = 0; i < h; i++) {
		for (j = 0; j < w; j++) {
			for (k = 0; k < c; k++) {
				*(dst_data + k * h * w + i * w + j) = *(src_data + i * w * c + j * c + k);
			}
		}
	}
}

inline static int32_t __clip(float val, float min, float max) {
    float f = val <= min ? min : (val >= max ? max : val);
    return f;
}

static void qnt_f32_to_affine(uint8_t *qnt, uint8_t zp, float scale, float *f32, int num) {
    float *src_ptr = f32;
    int i = 0;
    float dst_val = 0.0;

    for (; i < num; i++)
    {
        //dst_val = ((*src_ptr) / scale) + zp;
        dst_val = *src_ptr;
        *qnt = (uint8_t)dst_val;//__clip(dst_val, 0, 255);
        src_ptr++;
        qnt++;
    }
}

static float deqnt_affine_to_f32(uint8_t qnt, uint32_t zp, float scale)
{
    return ((float)qnt - (float)zp) * scale;
}

SiamRpnTracker::SiamRpnTracker() {
    log_i("SiamRpnTracker construct");
    track_instance_size = 288;
    track_exemplar_size = 128;
    track_base_size = 0;
    track_context_amount = 0.5;
    track_penalty_k = 0.16;
    track_window_influence = 0.4;
    track_lr = 0.3;
    anchor_stride = 8;
    anchor_ratios = { 0.33, 0.5, 1, 2, 3 };
    anchor_scales = { 8 };
    anchor_nump = 5;
    score_size = (track_instance_size - track_exemplar_size) / anchor_stride + 1 + track_base_size;

    FILE* template_fp = fopen(tem_model_path.c_str(), "rb");
    FILE* track_fp = fopen(track_model_path.c_str(), "rb");
    FILE* rpn_fp = fopen(rpn_model_path.c_str(), "rb");

    long template_fsize,track_fsize,rpn_fsize;

    fseek(template_fp, 0, SEEK_END);
    template_fsize = ftell(template_fp);
    template_modle = malloc(template_fsize);
    fseek(template_fp, 0, SEEK_SET);
    size_t template_len = fread((unsigned char*)template_modle,1,template_fsize,template_fp);

    fseek(track_fp, 0, SEEK_END);
    track_fsize = ftell(track_fp);
    track_modle = malloc(track_fsize);
    fseek(track_fp, 0, SEEK_SET);
    size_t track_len = fread((unsigned char*)track_modle,1,track_fsize,track_fp);

    fseek(rpn_fp, 0, SEEK_END);
    rpn_fsize = ftell(rpn_fp);
    rpn_modle = malloc(rpn_fsize);
    fseek(rpn_fp, 0, SEEK_SET);
    size_t rpn_len = fread((unsigned char*)rpn_modle,1,rpn_fsize,rpn_fp);

    int ret = rknn_init(&template_session,template_modle,template_fsize,0);
    log_i("template init res is %d", ret);
    log_i("template_fsize is %ld", template_fsize);
    log_i("template_len is %d", template_len);

    rknn_tensor_attr input_attrs;
    memset(&input_attrs, 0, sizeof(input_attrs));
    input_attrs.index = 0;
    rknn_query(template_session, RKNN_QUERY_INPUT_ATTR, &(input_attrs), sizeof(rknn_tensor_attr));
    printRKNNTensor(&(input_attrs));
    template_input_attrs = input_attrs;

    tem_inputs_mem = rknn_create_mem(template_session, template_input_attrs.size);
    rknn_set_io_mem(template_session, tem_inputs_mem, &template_input_attrs);
    log_i("virt_addr = %p, phys_addr = 0x%llx, tem_inputs_mem->fd:%d", tem_inputs_mem->logical_addr, tem_inputs_mem->physical_addr, tem_inputs_mem->fd);

    memset(&template_output_attrs, 0, sizeof(template_output_attrs));
    template_output_attrs.index = 0;
    rknn_query(template_session, RKNN_QUERY_OUTPUT_ATTR, &(template_output_attrs), sizeof(rknn_tensor_attr));
    printRKNNTensor(&(template_output_attrs));

    tem_outputs_mem = rknn_create_mem(template_session, template_output_attrs.size);
    rknn_set_io_mem(template_session, tem_outputs_mem, &template_output_attrs);
    log_i("virt_addr = %p, phys_addr = 0x%llx, tem_outputs_mem->fd:%d", tem_outputs_mem->logical_addr, tem_outputs_mem->physical_addr, tem_outputs_mem->fd);

    ret = rknn_init(&track_session,track_modle,track_fsize,0);
    log_i("track init res is %d", ret);
    log_i("track_fsize is %lld", track_fsize);
    log_i("track_len is %d", track_len);
    
    memset(&input_attrs, 0, sizeof(input_attrs));
    input_attrs.index = 0;
    rknn_query(track_session, RKNN_QUERY_INPUT_ATTR, &(input_attrs), sizeof(rknn_tensor_attr));
    printRKNNTensor(&(input_attrs));
    track_input_attrs = input_attrs;

    track_inputs_mem = rknn_create_mem(track_session, track_input_attrs.size);
    rknn_set_io_mem(track_session, track_inputs_mem, &track_input_attrs);
    log_i("virt_addr = %p, phys_addr = 0x%llx, track_inputs_mem->fd:%d", track_inputs_mem->logical_addr, track_inputs_mem->physical_addr, track_inputs_mem->fd);

    memset(&track_output_attrs, 0, sizeof(track_output_attrs));
    track_output_attrs.index = 0;
    rknn_query(track_session, RKNN_QUERY_OUTPUT_ATTR, &(track_output_attrs), sizeof(rknn_tensor_attr));
    printRKNNTensor(&(track_output_attrs));

    track_outputs_mem = rknn_create_mem(track_session, track_output_attrs.size);
    rknn_set_io_mem(track_session, track_outputs_mem, &track_output_attrs);
    log_i("virt_addr = %p, phys_addr = 0x%llx, track_outputs_mem->fd:%d", track_outputs_mem->logical_addr, track_outputs_mem->physical_addr, track_outputs_mem->fd);

    ret = rknn_init(&rpn_session,rpn_modle,rpn_fsize,0);
    log_i("rpn init res is %d", ret);
    log_i("rpn_fsize is %lld", rpn_fsize);
    log_i("rpn_len is %d", rpn_len);
    // printRKNNAttr(rpn_session);

    for (int i = 0; i < 2; i++) {
        memset(&rpn_input_attrs[i], 0, sizeof(rpn_input_attrs[i]));
        rpn_input_attrs[i].index = i;
        rknn_query(rpn_session, RKNN_QUERY_INPUT_ATTR, &(rpn_input_attrs[i]), sizeof(rknn_tensor_attr));
        printRKNNTensor(&(rpn_input_attrs[i]));

        memset(&rpn_output_attrs[i], 0, sizeof(rpn_output_attrs[i]));
        rpn_output_attrs[i].index = i;
        rknn_query(rpn_session, RKNN_QUERY_OUTPUT_ATTR, &(rpn_output_attrs[i]), sizeof(rknn_tensor_attr));
        printRKNNTensor(&(rpn_output_attrs[i]));

        rpn_outputs_mem[i] = rknn_create_mem(rpn_session, rpn_output_attrs[i].size);
        rknn_set_io_mem(rpn_session, rpn_outputs_mem[i], &rpn_output_attrs[i]);
        log_i("virt_addr = %p, phys_addr = 0x%llx, rpn_outputs_mem[%d]->fd:%d", rpn_outputs_mem[i]->logical_addr, rpn_outputs_mem[i]->physical_addr, i, rpn_outputs_mem[i]->fd);

        rknn_set_io_mem(rpn_session, rpn_outputs_mem[i], &rpn_output_attrs[i]);

        out_scales[i] = rpn_output_attrs[i].scale;
        out_zps[i] = rpn_output_attrs[i].zp;
    }

    rknn_set_io_mem(rpn_session, tem_outputs_mem, &rpn_input_attrs[0]);
    rknn_set_io_mem(rpn_session, track_outputs_mem, &rpn_input_attrs[1]);

    anchor = generate_anchors();
    xarray<double> hanning = xt::cos(xarray<double>(xt::arange(score_size) * (M_PI * 2. / static_cast<double>(score_size - 1)))) * -0.5 + 0.5;
    xarray<double> d_hanning = xt::empty<double>({ score_size, score_size });
    for (int r = 0; r < score_size; r++)
    {
        for (int c = 0; c < score_size; c++)
        {
            d_hanning.at(r, c) = hanning.at(r) * hanning.at(c);
        }
    }
    window = xt::tile(flatten(d_hanning), anchor_nump);
  
    fclose(template_fp);
    fclose(track_fp);
    fclose(rpn_fp);

    log_i("free model");
    xfree(template_modle);
    xfree(track_modle);
    xfree(rpn_modle);

    rknn_run(template_session, NULL);
    rknn_run(track_session, NULL);
    rknn_run(rpn_session, NULL);
}

SiamRpnTracker::~SiamRpnTracker() {
    rknn_destroy_mem(template_session, tem_inputs_mem);
    rknn_destroy_mem(template_session, tem_outputs_mem);
    rknn_destroy_mem(track_session, track_inputs_mem);
    rknn_destroy_mem(track_session, track_outputs_mem);
    rknn_destroy_mem(rpn_session, rpn_outputs_mem[0]);
    rknn_destroy_mem(rpn_session, rpn_outputs_mem[1]);

    rknn_destroy(template_session);
    rknn_destroy(track_session);
    rknn_destroy(rpn_session); 
}

int SiamRpnTracker::preprocess_img(MEDIA_BUFFER fram_mb, cv::Size frame_size ,Point2d K_center_pos) {
    // time_checker timer;
    // timer.start();
    double w_z = box_size.width + track_context_amount * (box_size.width + box_size.height);
    double h_z = box_size.height + track_context_amount * (box_size.width + box_size.height);
    double s_z = sqrt(w_z * h_z);
    double s_x = s_z * ((double)track_instance_size / (double)track_exemplar_size);
    // s_x = s_x < 200 ? 200.0 : s_x;

    center_pos = K_center_pos;
    // log_i("center_pos x = %f,y = %f s",center_pos.x,center_pos.y);

    int ret = get_subwindow(fram_mb, frame_size, center_pos, track_instance_size, round(s_x), track_inputs_mem);
    // timer.stop();
    // timer.show_distance("Update stage ---- get subwindow cost time");
    return ret;
}

track_result SiamRpnTracker::postprocess_img(cv::Size frame_size) {
    // timer.start();
    int whole_size = track_instance_size * track_instance_size * 3;

    double w_z = box_size.width + track_context_amount * (box_size.width + box_size.height);
    double h_z = box_size.height + track_context_amount * (box_size.width + box_size.height);
    double s_z = sqrt(w_z * h_z);
    double scale_z = track_exemplar_size / s_z;

    uint8_t* mask_head_outputs[2] = {(uint8_t *)rpn_outputs_mem[0]->logical_addr, (uint8_t *)rpn_outputs_mem[1]->logical_addr,};

    float max_pscore = 0.0;
    float max_score = 0.0;
    size_t max_index = 0;
    float max_penalty = 0.0;
    float max_box_w = 0.0;
    float max_box_h = 0.0;
    for (size_t i = 0; i < buffer_output_number; ++i) {
        float score1 = deqnt_affine_to_f32(*(mask_head_outputs[0] + i), out_zps[0], out_scales[0]);
        float score2 = deqnt_affine_to_f32(*(mask_head_outputs[0] + i + buffer_output_number), out_zps[0], out_scales[0]);
        float score2_exp = std::exp(score2);
        float score_softmax = score2_exp / (score2_exp + std::exp(score1));
        float pred_bbox_w =
            std::exp(deqnt_affine_to_f32(*(mask_head_outputs[1] + i + buffer_output_number * 2), out_zps[1], out_scales[1])) *
            anchor.at(i, 2);
        float pred_bbox_H =
            std::exp(deqnt_affine_to_f32(*(mask_head_outputs[1] + i + buffer_output_number * 3), out_zps[1], out_scales[1])) *
            anchor.at(i, 3);
        float s_c =
            change_t(sz_t(pred_bbox_w, pred_bbox_H) /
                        sz_t(float(box_size.width * scale_z), float(box_size.height * scale_z)));
        float r_c = change_t(float(box_size.width / box_size.height) /
                                float(pred_bbox_w / pred_bbox_H));
        float penalty = std::exp(-(r_c * s_c - 1) * track_penalty_k);
        float pscore = penalty * score_softmax;
        pscore = pscore * (1 - track_window_influence) +
                    window[i] * track_window_influence;
        if (pscore > max_pscore) {
            max_pscore = pscore;
            max_score = score_softmax;
            max_index = i;
            max_penalty = penalty;
            max_box_w = pred_bbox_w;
            max_box_h = pred_bbox_H;
        }
    }
    float bbox_max_score[4] = {
        (deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index), out_zps[1], out_scales[1]) * anchor.at(max_index, 2) + anchor.at(max_index, 0)) / (float)scale_z,
        (deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index + buffer_output_number), out_zps[1], out_scales[1]) * anchor.at(max_index, 3) + anchor.at(max_index, 1)) / (float)scale_z,
        max_box_w / (float)scale_z, 
        max_box_h / (float)scale_z};
    float lr = max_penalty * max_score * track_lr;
    double cx = bbox_max_score[0] + center_pos.x;
    double cy = bbox_max_score[1] + center_pos.y;
    double width = box_size.width * (1 - lr) + bbox_max_score[2] * lr;
    double height = box_size.height * (1 - lr) + bbox_max_score[3] * lr;
    Rect2d box = bbox_clip(Rect2d(cx, cy, width, height), frame_size);
    center_pos = Point2d(cx, cy);
    box_size = Size2d(width, height);
    box = Rect2d(box.x - box.width / 2, box.y - box.height / 2, width, height);
    float bscore = max_score;
    // std::cout << "new result, box : " << box << ", bscore : " << bscore << std::endl; 
    track_result track_res;
    track_res.bbox = box;
    track_res.bscore = bscore;
    track_res.center_pos = center_pos;

    // timer.stop();
    // timer.show_distance("Update stage ---- postprocess cost time");

    return track_res;
}


track_result SiamRpnTracker::postprocess_img_without_smooth(cv::Size frame_size) {
    // timer.start();
    int whole_size = track_instance_size * track_instance_size * 3;

    double w_z = box_size.width + track_context_amount * (box_size.width + box_size.height);
    double h_z = box_size.height + track_context_amount * (box_size.width + box_size.height);
    double s_z = sqrt(w_z * h_z);
    double scale_z = track_exemplar_size / s_z;

    uint8_t* mask_head_outputs[2] = {(uint8_t *)rpn_outputs_mem[0]->logical_addr, (uint8_t *)rpn_outputs_mem[1]->logical_addr,};

    float max_score = 0.0;
    size_t max_index = 0;
    for (size_t i = 0; i < buffer_output_number; ++i) {
        float score1 = deqnt_affine_to_f32(*(mask_head_outputs[0] + i), out_zps[0], out_scales[0]);
        float score2 = deqnt_affine_to_f32(*(mask_head_outputs[0] + i + buffer_output_number), out_zps[0], out_scales[0]);
        float score2_exp = std::exp(score2);
        float score_softmax = score2_exp / (score2_exp + std::exp(score1));

        if (score_softmax > max_score) {
            max_score = score_softmax;
            max_index = i;
        }
    }

    float bbox_max_score[4] = {
        (deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index), out_zps[1], out_scales[1]) * anchor.at(max_index, 2) + anchor.at(max_index, 0)) / (float)scale_z,
        (deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index + buffer_output_number), out_zps[1], out_scales[1]) * anchor.at(max_index, 3) + anchor.at(max_index, 1)) / (float)scale_z,
        std::exp(deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index + buffer_output_number * 2), out_zps[1], out_scales[1])) * anchor.at(max_index, 2) / (float)scale_z, 
        std::exp(deqnt_affine_to_f32(*(mask_head_outputs[1] + max_index + buffer_output_number * 3), out_zps[1], out_scales[1])) * anchor.at(max_index, 3) / (float)scale_z};

    double cx = bbox_max_score[0] + center_pos.x;
    double cy = bbox_max_score[1] + center_pos.y;
    double width = bbox_max_score[2];
    double height = bbox_max_score[3];

    Rect2d box = bbox_clip(Rect2d(cx, cy, bbox_max_score[2], bbox_max_score[3]), frame_size);
    center_pos = Point2d(cx, cy);

    box = Rect2d(box.x - box.width / 2, box.y - box.height / 2, width, height);
    float bscore = max_score;
    // std::cout << "new result, box : " << box << ", bscore : " << bscore << std::endl; 
    track_result track_res;
    track_res.bbox = box;
    track_res.bscore = bscore;
    track_res.center_pos = center_pos;

    // timer.stop();
    // timer.show_distance("Update stage ---- postprocess cost time");

    return track_res;
}

track_result SiamRpnTracker::track(cv::Size frame_size, bool redetect) {
    // time_checker timer;

    // timer.start();

    rknn_run(track_session,NULL);

    rknn_run(rpn_session,NULL);

    // timer.stop();
    // timer.show_distance("Update stage ---- NPU cost time");

    track_result track_res;
    if (redetect) {
        track_res = postprocess_img_without_smooth(frame_size);
    }
    else {
        track_res = postprocess_img(frame_size);
    }
    
    return track_res;
}

int SiamRpnTracker::init(MEDIA_BUFFER fram_mb, cv::Size frame_size, Rect2d box) {
    center_pos = box.tl() + Point2d((box.width - 1) / 2, (box.height - 1) / 2);
    box_size = box.size();
    int whole_size = track_exemplar_size * track_exemplar_size * 3;

    double w_z = box_size.width + track_context_amount * (box.width + box.height);
    double h_z = box_size.height + track_context_amount * (box.width + box.height);
    double s_z = round(sqrt(w_z * h_z));
    
    int ret = get_subwindow(fram_mb, frame_size, center_pos, track_exemplar_size, s_z, tem_inputs_mem);
    if(ret < 0)
        return 0;

    rknn_run(template_session, 0);

    return 1;
}

int SiamRpnTracker::initWithBGR(MEDIA_BUFFER fram_mb, cv::Size frame_size, Rect2d box) {
    center_pos = box.tl() + Point2d((box.width - 1) / 2, (box.height - 1) / 2);
    box_size = box.size();
    int whole_size = track_exemplar_size * track_exemplar_size * 3;

    double w_z = box_size.width + track_context_amount * (box.width + box.height);
    double h_z = box_size.height + track_context_amount * (box.width + box.height);
    double s_z = round(sqrt(w_z * h_z));
    
    int ret = get_subwindowWithBGR(fram_mb, frame_size, center_pos, track_exemplar_size, s_z, tem_inputs_mem);
    if(ret < 0)
        return 0;

    rknn_run(template_session, 0);

    return 1;
}

xarray<float> SiamRpnTracker::generate_anchors() {
    xarray<float> anchor = xt::zeros<float>({ anchor_nump, 4 });
    
    int size = anchor_stride * anchor_stride;
    int count = 0;
    for (int ri = 0; ri < anchor_ratios.size(); ri++) {
        float r = anchor_ratios[ri];
        int ws = sqrt(size / r);
        int hs = ws * r;

        for (int si = 0; si < anchor_scales.size(); si++) {
            float s = anchor_scales[si];
            float w = ws * s;
            float h = hs * s;
            xt::view(anchor, count, xt::all()) = xarray<float>({-w * 0.5f, -h * 0.5f, w * 0.5f, h * 0.5f});
            count++;
        }
    }

    xarray<int> xx, yy;
    xarray<float> x1 = xt::view(anchor, xt::all(), 0);
    xarray<float> y1 = xt::view(anchor, xt::all(), 1);
    xarray<float> x2 = xt::view(anchor, xt::all(), 2);
    xarray<float> y2 = xt::view(anchor, xt::all(), 3);

    anchor = xt::stack(xtuple((x1 + x2) * 0.5f, (y1 + y2) * 0.5f, x2 - x1, y2 - y1), 1);
    anchor = xt::tile(anchor, { 1,score_size * score_size });
    anchor = anchor.reshape({ -1,4 });

    int ori = -(score_size / 2) * anchor_stride;
    
    xarray<int> me, ti;
    vector<int> mea, tia;

    for (int dx = 0; dx < score_size; dx++) {
        mea.emplace_back(ori + anchor_stride * dx);
    }
    me = xt::adapt(mea, { (int)mea.size()});

    for (int dy = 0; dy < score_size; dy++) {
        tia.emplace_back(ori + anchor_stride * dy);
    }
    ti = xt::adapt(tia, { (int)tia.size() });

    auto meshres = meshgrid(me, ti);
    xx = get<0>(meshres);
    yy = get<1>(meshres);
    xx = xt::transpose(xx, { 1,0 });
    yy = xt::transpose(yy, { 1,0 });

    xx = flatten(xx);
    yy = flatten(yy);
    xx = tile(xx.reshape({ 1,xx.size() }), { anchor_nump, 1 });
    yy = tile(yy.reshape({ 1,yy.size() }), { anchor_nump, 1 });
    xx = flatten(xx);
    yy = flatten(yy);

    xt::view(anchor, xt::all(), 0) = xt::cast<float>(xx);
    xt::view(anchor, xt::all(), 1) = xt::cast<float>(yy);

    return anchor;
}

xarray<float> SiamRpnTracker::convert_score(xarray<float> score) {
    xarray<float> scf;
    scf = xarray<float>(xt::transpose(score, { 1,2,3,0 }));
    scf = xarray<float>(xt::transpose(scf.reshape({ 2,-1 }), { 1,0 }));
    scf = xt::view(softmax(scf, 1), xt::all(), 1);
    return scf;
}

xarray<float> SiamRpnTracker::convert_box(xarray<float> delta, xarray<float> anchor) {
    xarray<float> bod;
    bod = xt::transpose(delta, { 1,2,3,0 });
    bod = bod.reshape({ 4,-1 });

    xt::view(bod, 0, xt::all()) = xt::view(bod, 0, xt::all()) * xt::view(anchor, xt::all(), 2) + xt::view(anchor, xt::all(), 0);
    xt::view(bod, 1, xt::all()) = xt::view(bod, 1, xt::all()) * xt::view(anchor, xt::all(), 3) + xt::view(anchor, xt::all(), 1);
    xt::view(bod, 2, xt::all()) = xt::exp(xt::view(bod, 2, xt::all())) * xt::view(anchor, xt::all(), 2);
    xt::view(bod, 3, xt::all()) = xt::exp(xt::view(bod, 3, xt::all())) * xt::view(anchor, xt::all(), 3);

    return bod;
}

Rect2d SiamRpnTracker::bbox_clip(Rect2d box, Size boundary) {
    Rect2d resbox;
    resbox.x = max((double)0, min(box.x, (double)boundary.width));
    resbox.y = max((double)0, min(box.y, (double)boundary.height));
    resbox.width = max((double)10, min(box.width, (double)boundary.width));
    resbox.height = max((double)10, min(box.height, (double)boundary.height));
    return resbox;
}

int SiamRpnTracker::get_subwindow(MEDIA_BUFFER fram_mb, cv::Size frame_size, Point2d center_pos, int model_size, int original_size, rknn_tensor_mem* tensor_mem) {
    time_checker timer;
    // timer.start();
    static int idx = 0;
    int r = frame_size.height, c = frame_size.width, k = 3;
    int whole_size = model_size * model_size * k;
    unsigned char *im_patch_ptr;
    xarray<uchar> im_patch;

    unsigned char* crop_buf = NULL;
    unsigned char* resize_buf = NULL;
    unsigned char* transpose_buf = NULL;
    im_rect 		src_rect;
    im_rect 		crop_rect;
    im_rect         prect;
    rga_buffer_t 	src;
    rga_buffer_t 	crop;
    rga_buffer_t    pat;
    MEDIA_BUFFER src_mb = fram_mb;
    MEDIA_BUFFER crop_mb = NULL;

    float scale = (float)model_size / (float)original_size;
    float s_z = original_size;
    float cc = (s_z + 1) / 2;
    
    float context_xmin = floor(center_pos.x - cc + 0.5);
    float context_xmax = context_xmin + s_z - 1;
    float context_ymin = floor(center_pos.y - cc + 0.5);
    float context_ymax = context_ymin + s_z - 1;

    int valid_context_xmin = ((int)context_xmin < 0) ? 0 : (int)context_xmin;
    int valid_context_xmax = ((int)context_xmax < c - 1) ? (int)context_xmax : (c - 1);
    int valid_context_ymin = ((int)context_ymin < 0) ? 0 : (int)context_ymin;
    int valid_context_ymax = ((int)context_ymax < r - 1) ? (int)context_ymax : (r - 1);

    int left_pad = std::max(0.f, -context_xmin);
    int top_pad = std::max(0.f, -context_ymin);
    int right_pad = std::max(0.f, context_xmax - frame_size.width + 1);
    int bottom_pad = std::max(0.f, context_ymax - frame_size.height + 1);

    context_xmin += left_pad;
    context_xmax += left_pad;
    context_ymin += top_pad;
    context_ymax += top_pad;

    memset(&src_rect, 0, sizeof(src_rect));
	memset(&crop_rect, 0, sizeof(crop_rect));
    memset(&prect, 0, sizeof(prect));
	memset(&src, 0, sizeof(src));
	memset(&crop, 0, sizeof(crop));
    memset(&pat, 0, sizeof(pat));

    int context_width = valid_context_xmax - valid_context_xmin + 1, context_height = valid_context_ymax - valid_context_ymin + 1;
	// printf("context_width = %d, context_height = %d, original_size = %d.\r\n", context_width, context_height, original_size);

    memset(tensor_mem->logical_addr, 0, tensor_mem->size);// 每次要清空缓冲区，因为RGA的结果只会覆盖一部分图像，其余地方会保留上次数据

    // 必须同步数据，否则会出现黑条
    dma_sync_cpu_to_device(RK_MPI_MB_GetFD(src_mb));
    dma_sync_cpu_to_device(tensor_mem->fd);

    src = wrapbuffer_fd(RK_MPI_MB_GetFD(src_mb), frame_size.width, frame_size.height, RK_FORMAT_YCbCr_420_SP);
    crop = wrapbuffer_fd(tensor_mem->fd, model_size, model_size, RK_FORMAT_BGR_888); //RGA对RGB888要4字节对齐

    //RGA对yuv要2字节对齐
    if (valid_context_xmin % 2 != 0)
        valid_context_xmin++;
    if (valid_context_ymin % 2 != 0)
        valid_context_ymin++;

    if (context_width % 2 != 0)
        context_width--;
    if (context_height % 2 != 0)
        context_height--;

    src_rect = {valid_context_xmin, valid_context_ymin, context_width, context_height};
    crop_rect = {int(left_pad * scale), int(top_pad * scale), int(context_width * scale), int(context_height * scale)};
    
    // printf("valid_context_xmin %d, valid_context_ymin %d, context_width %d, context_height %d, scale %f, left_pad*scale %f,top_pad*scale %f,context_width*scale %f, context_height*scale %f\n",valid_context_xmin, valid_context_ymin, context_width, context_height, scale, left_pad*scale,top_pad*scale,context_width*scale, context_height*scale);
    // printf("crop_rect.x %d, crop_rect.y %d, crop_rect.width %d, crop_rect.height %d\n", crop_rect.x, crop_rect.y, crop_rect.width, crop_rect.height);

    IM_STATUS STATUS = improcess(src, crop, pat, src_rect, crop_rect, prect, IM_SYNC);
    if (STATUS != IM_STATUS_SUCCESS) {
		log_e("improcess failed: %s", imStrError(STATUS));
        return -1;
	}

    // 必须同步数据，否则会出现黑条
    dma_sync_device_to_cpu(RK_MPI_MB_GetFD(src_mb));
    dma_sync_device_to_cpu(tensor_mem->fd);

    // // 创建一个Mat对象，并指定图像尺寸和数据类型
    // static int img_count = 0;
    // cv::Mat image(model_size, model_size, CV_8UC3,  (uint8_t*)tensor_mem->logical_addr);

    // // 保存图像到文件
    // std::string filename = std::string("/userdata/install/baike_test/track_test/").append(std::to_string(img_count)).append(std::string(".jpg")); 
    // cv::imwrite(filename, image);
    
    // img_count++;
    return 0;
}


int SiamRpnTracker::get_subwindowWithBGR(MEDIA_BUFFER fram_mb, cv::Size frame_size, Point2d center_pos, int model_size, int original_size, rknn_tensor_mem* tensor_mem) {
    time_checker timer;
    // timer.start();
    static int idx = 0;
    int r = frame_size.height, c = frame_size.width, k = 3;
    int whole_size = model_size * model_size * k;
    unsigned char *im_patch_ptr;
    xarray<uchar> im_patch;

    unsigned char* crop_buf = NULL;
    unsigned char* resize_buf = NULL;
    unsigned char* transpose_buf = NULL;
    im_rect 		src_rect;
    im_rect 		crop_rect;
    im_rect         prect;
    rga_buffer_t 	src;
    rga_buffer_t 	crop;
    rga_buffer_t    pat;
    MEDIA_BUFFER src_mb = fram_mb;
    MEDIA_BUFFER crop_mb = NULL;

    float scale = (float)model_size / (float)original_size;
    float s_z = original_size;
    float cc = (s_z + 1) / 2;
    
    float context_xmin = floor(center_pos.x - cc + 0.5);
    float context_xmax = context_xmin + s_z - 1;
    float context_ymin = floor(center_pos.y - cc + 0.5);
    float context_ymax = context_ymin + s_z - 1;

    int valid_context_xmin = ((int)context_xmin < 0) ? 0 : (int)context_xmin;
    int valid_context_xmax = ((int)context_xmax < c - 1) ? (int)context_xmax : (c - 1);
    int valid_context_ymin = ((int)context_ymin < 0) ? 0 : (int)context_ymin;
    int valid_context_ymax = ((int)context_ymax < r - 1) ? (int)context_ymax : (r - 1);

    int left_pad = std::max(0.f, -context_xmin);
    int top_pad = std::max(0.f, -context_ymin);
    int right_pad = std::max(0.f, context_xmax - frame_size.width + 1);
    int bottom_pad = std::max(0.f, context_ymax - frame_size.height + 1);

    context_xmin += left_pad;
    context_xmax += left_pad;
    context_ymin += top_pad;
    context_ymax += top_pad;

    memset(&src_rect, 0, sizeof(src_rect));
	memset(&crop_rect, 0, sizeof(crop_rect));
    memset(&prect, 0, sizeof(prect));
	memset(&src, 0, sizeof(src));
	memset(&crop, 0, sizeof(crop));
    memset(&pat, 0, sizeof(pat));

    int context_width = valid_context_xmax - valid_context_xmin + 1, context_height = valid_context_ymax - valid_context_ymin + 1;
	// printf("context_width = %d, context_height = %d, original_size = %d.\r\n", context_width, context_height, original_size);

    memset(tensor_mem->logical_addr, 0, tensor_mem->size);// 每次要清空缓冲区，因为RGA的结果只会覆盖一部分图像，其余地方会保留上次数据

    // 必须同步数据，否则会出现黑条
    dma_sync_cpu_to_device(RK_MPI_MB_GetFD(src_mb));
    dma_sync_cpu_to_device(tensor_mem->fd);

    src = wrapbuffer_fd(RK_MPI_MB_GetFD(src_mb), frame_size.width, frame_size.height, RK_FORMAT_BGR_888);
    crop = wrapbuffer_fd(tensor_mem->fd, model_size, model_size, RK_FORMAT_BGR_888); //RGA对RGB888要4字节对齐

    //RGA对yuv要2字节对齐
    if (valid_context_xmin % 2 != 0)
        valid_context_xmin++;
    if (valid_context_ymin % 2 != 0)
        valid_context_ymin++;

    if (context_width % 2 != 0)
        context_width--;
    if (context_height % 2 != 0)
        context_height--;

    src_rect = {valid_context_xmin, valid_context_ymin, context_width, context_height};
    crop_rect = {int(left_pad * scale), int(top_pad * scale), int(context_width * scale), int(context_height * scale)};
    
    // printf("valid_context_xmin %d, valid_context_ymin %d, context_width %d, context_height %d, scale %f, left_pad*scale %f,top_pad*scale %f,context_width*scale %f, context_height*scale %f\n",valid_context_xmin, valid_context_ymin, context_width, context_height, scale, left_pad*scale,top_pad*scale,context_width*scale, context_height*scale);
    // printf("crop_rect.x %d, crop_rect.y %d, crop_rect.width %d, crop_rect.height %d\n", crop_rect.x, crop_rect.y, crop_rect.width, crop_rect.height);

    IM_STATUS STATUS = improcess(src, crop, pat, src_rect, crop_rect, prect, IM_SYNC);
    if (STATUS != IM_STATUS_SUCCESS) {
		log_e("improcess failed: %s", imStrError(STATUS));
        return -1;
	}

    // 必须同步数据，否则会出现黑条
    dma_sync_device_to_cpu(RK_MPI_MB_GetFD(src_mb));
    dma_sync_device_to_cpu(tensor_mem->fd);

    // // 创建一个Mat对象，并指定图像尺寸和数据类型
    // static int img_count = 0;
    // cv::Mat image(model_size, model_size, CV_8UC3,  (uint8_t*)tensor_mem->logical_addr);

    // // 保存图像到文件
    // std::string filename = std::string("/userdata/install/baike_test/track_test/bgr_").append(std::to_string(img_count)).append(std::string(".jpg")); 
    // cv::imwrite(filename, image);
    
    // img_count++;
    return 0;
}

void  SiamRpnTracker::printRKNNTensor(rknn_tensor_attr *attr) {
    log_i("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
            "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f",
            attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
            attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
            attr->qnt_type, attr->fl, attr->zp, attr->scale);
}

int SiamRpnTracker::printRKNNAttr(rknn_context ctx) {
    rknn_input_output_num io_num;
    int ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0)
    {
        log_e("rknn_init error ret=%d", ret);
        return -1;
    }
    int input_num = io_num.n_input;
    int output_num = io_num.n_output;
    log_i("model input num: %d, output num: %d", input_num,
           output_num);

    rknn_tensor_attr input_attrs[input_num];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < input_num; i++)
    {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]),
                            sizeof(rknn_tensor_attr));
        printRKNNTensor(&(input_attrs[i]));
    }

    rknn_tensor_attr output_attrs[output_num];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int k = 0; k < output_num; k++)
    {
        output_attrs[k].index = k;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[k]),
                            sizeof(rknn_tensor_attr));
        printRKNNTensor(&(output_attrs[k]));
    }
    return 0;
}
